import os
import paddle
"""
Export a YOLOv5 PyTorch model to TorchScript, ONNX, CoreML, TensorFlow (saved_model, pb, TFLite, TF.js,) formats
TensorFlow exports authored by https://github.com/zldrobit

Usage:
    $ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml saved_model pb tflite tfjs

Inference:
    $ python path/to/detect.py --weights yolov5s.pt
                                         yolov5s.onnx  (must export with --dynamic)
                                         yolov5s_saved_model
                                         yolov5s.pb
                                         yolov5s.tflite

TensorFlow.js:
    $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
    $ npm install
    $ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model
    $ npm start
"""
import argparse
import subprocess
import sys
import time
from pathlib import Path
import pandas as pd
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))
from models.common import Conv
from models.experimental import attempt_load
from models.yolo import Detect
from utils.activations import SiLU
from utils.datasets import LoadImages
from utils.general import colorstr, check_dataset, check_img_size, check_requirements, file_size, print_args, set_logging, url2file
from utils.torch_utils import select_device


def export_formats():
    x = [['PyTorch', '-', '.pt'], ['TorchScript', 'torchscript',
        '.torchscript'], ['ONNX', 'onnx', '.onnx'], ['OpenVINO', 'openvino',
        '_openvino_model'], ['TensorRT', 'engine', '.engine'], ['CoreML',
        'coreml', '.mlmodel'], ['TensorFlow SavedModel', 'saved_model',
        '_saved_model'], ['TensorFlow GraphDef', 'pb', '.pb'], [
        'TensorFlow Lite', 'tflite', '.tflite'], ['TensorFlow Edge TPU',
        'edgetpu', '_edgetpu.tflite'], ['TensorFlow.js', 'tfjs', '_web_model']]
    return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix'])


def export_torchscript(model, im, file, optimize, prefix=colorstr(
    'TorchScript:')):
    try:
        print(f'\n{prefix} starting export with torch {paddle.__version__}...')
        f = file.with_suffix('.torchscript.pt')
        ts = paddle.jit.trace(model, im, strict=False)
        (optimize_for_mobile(ts) if optimize else ts).save(f)
        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
    except Exception as e:
        print(f'{prefix} export failure: {e}')


def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=
    colorstr('ONNX:')):
    try:
        check_requirements(('onnx',))
        import onnx
        print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
        f = file.with_suffix('.onnx')
        paddle.onnx.export(model, im, f, verbose=False, opset_version=opset,
            training=paddle.onnx.TrainingMode.TRAINING if train else paddle.
            onnx.TrainingMode.EVAL, do_constant_folding=not train,
            input_names=['images'], output_names=['output'], dynamic_axes={
            'images': {(0): 'batch', (2): 'height', (3): 'width'}, 'output':
            {(0): 'batch', (1): 'anchors'}} if dynamic else None)
        model_onnx = onnx.load(f)
        onnx.checker.check_model(model_onnx)
        if simplify:
            try:
                check_requirements(('onnx-simplifier',))
                import onnxsim
                print(
                    f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...'
                    )
                model_onnx, check = onnxsim.simplify(model_onnx,
                    dynamic_input_shape=dynamic, input_shapes={'images':
                    list(tuple(im.shape))} if dynamic else None)
                assert check, 'assert check failed'
                onnx.save(model_onnx, f)
            except Exception as e:
                print(f'{prefix} simplifier failure: {e}')
        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
        print(
            f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'"
            )
    except Exception as e:
        print(f'{prefix} export failure: {e}')


def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
    ct_model = None
    try:
        check_requirements(('coremltools',))
        import coremltools as ct
        print(
            f'\n{prefix} starting export with coremltools {ct.__version__}...')
        f = file.with_suffix('.mlmodel')
        model.train()
        ts = paddle.jit.trace(model, im, strict=False)
        ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=tuple
            (im.shape), scale=1 / 255.0, bias=[0, 0, 0])])
        ct_model.save(f)
        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
    except Exception as e:
        print(f'\n{prefix} export failure: {e}')
    return ct_model


def export_saved_model(model, im, file, dynamic, tf_nms=False, agnostic_nms
    =False, topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=
    0.25, prefix=colorstr('TensorFlow saved_model:')):
    keras_model = None
    try:
        import tensorflow as tf
        from tensorflow import keras
        from models.tf import TFModel, TFDetect
        print(f'\n{prefix} starting export with tensorflow {tf.__version__}...'
            )
        f = str(file).replace('.pt', '_saved_model')
        batch_size, ch, *imgsz = list(tuple(im.shape))
        tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=
            imgsz)
        im = tf.zeros((batch_size, *imgsz, 3))
        y = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class,
            topk_all, iou_thres, conf_thres)
        inputs = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else
            batch_size)
        outputs = tf_model.predict(inputs, tf_nms, agnostic_nms,
            topk_per_class, topk_all, iou_thres, conf_thres)
        keras_model = keras.Model(inputs=inputs, outputs=outputs)
        keras_model.trainable = False
        keras_model.summary()
        keras_model.save(f, save_format='tf')
        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
    except Exception as e:
        print(f'\n{prefix} export failure: {e}')
    return keras_model


def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
    try:
        import tensorflow as tf
        from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
        print(f'\n{prefix} starting export with tensorflow {tf.__version__}...'
            )
        f = file.with_suffix('.pb')
        m = tf.function(lambda x: keras_model(x))
        m = m.get_concrete_function(tf.TensorSpec(tuple(keras_model.inputs[
            0].shape), keras_model.inputs[0].dtype))
        frozen_func = convert_variables_to_constants_v2(m)
        frozen_func.graph.as_graph_def()
        tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(
            f.parent), name=f.name, as_text=False)
        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
    except Exception as e:
        print(f'\n{prefix} export failure: {e}')


def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=
    colorstr('TensorFlow Lite:')):
    try:
        import tensorflow as tf
        from models.tf import representative_dataset_gen
        print(f'\n{prefix} starting export with tensorflow {tf.__version__}...'
            )
        batch_size, ch, *imgsz = list(tuple(im.shape))
        f = str(file).replace('.pt', '-fp16.tflite')
        converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
        converter.target_spec.supported_types = [tf.float16]
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
        if int8:
            dataset = LoadImages(check_dataset(data)['train'], img_size=
                imgsz, auto=False)
            converter.representative_dataset = (lambda :
                representative_dataset_gen(dataset, ncalib))
            converter.target_spec.supported_ops = [tf.lite.OpsSet.
                TFLITE_BUILTINS_INT8]
            converter.target_spec.supported_types = []
            converter.inference_input_type = tf.uint8
            converter.inference_output_type = tf.uint8
            converter.experimental_new_quantizer = False
            f = str(file).replace('.pt', '-int8.tflite')
        tflite_model = converter.convert()
        open(f, 'wb').write(tflite_model)
        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
    except Exception as e:
        print(f'\n{prefix} export failure: {e}')


def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
    try:
        check_requirements(('tensorflowjs',))
        import re
        import tensorflowjs as tfjs
        print(
            f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...'
            )
        f = str(file).replace('.pt', '_web_model')
        f_pb = file.with_suffix('.pb')
        f_json = f + '/model.json'
        cmd = (
            f"tensorflowjs_converter --input_format=tf_frozen_model --output_node_names='Identity,Identity_1,Identity_2,Identity_3' {f_pb} {f}"
            )
        subprocess.run(cmd, shell=True)
        json = open(f_json).read()
        with open(f_json, 'w') as j:
            subst = re.sub(
                '{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, "Identity.?.?": {"name": "Identity.?.?"}, "Identity.?.?": {"name": "Identity.?.?"}, "Identity.?.?": {"name": "Identity.?.?"}}}'
                ,
                '{"outputs": {"Identity": {"name": "Identity"}, "Identity_1": {"name": "Identity_1"}, "Identity_2": {"name": "Identity_2"}, "Identity_3": {"name": "Identity_3"}}}'
                , json)
            j.write(subst)
        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
    except Exception as e:
        print(f'\n{prefix} export failure: {e}')


@paddle.no_grad()
def run(data=ROOT / 'data/coco128.yaml', weights=ROOT / 'yolov5s.pt', imgsz
    =(640, 640), batch_size=1, device='cpu', include=('torchscript', 'onnx',
    'coreml'), half=False, inplace=False, train=False, optimize=False, int8
    =False, dynamic=False, simplify=False, opset=12, topk_per_class=100,
    topk_all=100, iou_thres=0.45, conf_thres=0.25):
    t = time.time()
    include = [x.lower() for x in include]
    tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite',
        'tfjs'))
    imgsz *= 2 if len(imgsz) == 1 else 1
    file = Path(url2file(weights) if str(weights).startswith(('http:/',
        'https:/')) else weights)
    device = select_device(device)
    assert not (device.type == 'cpu' and half
        ), '--half only compatible with GPU export, i.e. use --device 0'
    model = attempt_load(weights, map_location=device, inplace=True, fuse=True)
    """Class Attribute: torch.Tensor.names, can not convert, please check whether it is torch.Tensor.*/torch.autograd.function.FunctionCtx.*/torch.distributions.Distribution.* and convert manually"""
    nc, names = model.nc, model.names
    gs = int(max(model.stride))
    imgsz = [check_img_size(x, gs) for x in imgsz]
    im = paddle.zeros(shape=[batch_size, 3, *imgsz]).to(device)
    if half:
        im, model = im.astype(dtype='float16'), model.astype(dtype='float16')
    model.train() if train else model.eval()
    for k, m in model.named_sublayers():
        if isinstance(m, Conv):
            if isinstance(m.act, paddle.nn.Silu):
                m.act = SiLU()
        elif isinstance(m, Detect):
            m.inplace = inplace
            m.onnx_dynamic = dynamic
    for _ in range(2):
        y = model(im)
    print(
        f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)"
        )
    if 'torchscript' in include:
        export_torchscript(model, im, file, optimize)
    if 'onnx' in include:
        export_onnx(model, im, file, opset, train, dynamic, simplify)
    if 'coreml' in include:
        export_coreml(model, im, file)
    if any(tf_exports):
        pb, tflite, tfjs = tf_exports[1:]
        assert not (tflite and tfjs
            ), 'TFLite and TF.js models must be exported separately, please pass only one type.'
        model = export_saved_model(model, im, file, dynamic, tf_nms=tfjs,
            agnostic_nms=tfjs, topk_per_class=topk_per_class, topk_all=
            topk_all, conf_thres=conf_thres, iou_thres=iou_thres)
        if pb or tfjs:
            export_pb(model, im, file)
        if tflite:
            export_tflite(model, im, file, int8=int8, data=data, ncalib=100)
        if tfjs:
            export_tfjs(model, im, file)
    print(
        f"""
Export complete ({time.time() - t:.2f}s)
Results saved to {colorstr('bold', file.parent.resolve())}
Visualize with https://netron.app"""
        )


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default=ROOT /
        'data/coco128.yaml', help='dataset.yaml path')
    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt',
        help='weights path')
    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=
        int, default=[640, 640], help='image (h, w)')
    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
    parser.add_argument('--device', default='cpu', help=
        'cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--half', action='store_true', help=
        'FP16 half-precision export')
    parser.add_argument('--inplace', action='store_true', help=
        'set YOLOv5 Detect() inplace=True')
    parser.add_argument('--train', action='store_true', help=
        'model.train() mode')
    parser.add_argument('--optimize', action='store_true', help=
        'TorchScript: optimize for mobile')
    parser.add_argument('--int8', action='store_true', help=
        'CoreML/TF INT8 quantization')
    parser.add_argument('--dynamic', action='store_true', help=
        'ONNX/TF: dynamic axes')
    parser.add_argument('--simplify', action='store_true', help=
        'ONNX: simplify model')
    parser.add_argument('--opset', type=int, default=13, help=
        'ONNX: opset version')
    parser.add_argument('--topk-per-class', type=int, default=100, help=
        'TF.js NMS: topk per class to keep')
    parser.add_argument('--topk-all', type=int, default=100, help=
        'TF.js NMS: topk for all classes to keep')
    parser.add_argument('--iou-thres', type=float, default=0.45, help=
        'TF.js NMS: IoU threshold')
    parser.add_argument('--conf-thres', type=float, default=0.25, help=
        'TF.js NMS: confidence threshold')
    parser.add_argument('--include', nargs='+', default=['torchscript',
        'onnx'], help=
        'available formats are (torchscript, onnx, coreml, saved_model, pb, tflite, tfjs)'
        )
    opt = parser.parse_args()
    print_args(FILE.stem, opt)
    return opt


def main(opt):
    set_logging()
    run(**vars(opt))


if __name__ == '__main__':
    opt = parse_opt()
    main(opt)
